from functools import singledispatch

from number_theory.gcd import gcd as gcd_number
from number_theory.gcd import inv_mod as inv_mod_number
from number_theory.gcd import lcm_list
from polynomial.polynomial_Q import PolynomialQ
from polynomial.polynomial_Z import PolynomialZ
from polynomial.polynomial_Zn import PolynomialZn
from prime.prime import Prime

__all__ = ["gcd", "exgcd", "inv_mod"]



@singledispatch
def gcd_template(m, n):
    while m:
        m, n = n % m, m
    return n.monic()


@gcd_template.register(PolynomialZ)
def _gcd_polynomial_z(m, n):
    m = m.primitive()
    n = n.primitive()
    if not m:
        return n
    elif not n:
        return m
    x, y = m.leading_coeff, n.leading_coeff
    d = gcd_number(x, y)
    e = x // d * y

    result = PolynomialZ()
    product = 1
    for prime in filter(lambda p: e % p, Prime.generate()):
        PolyType = PolynomialZn(prime)
        g = gcd(PolyType(m), PolyType(n))
        if g.deg == 0:
            return PolynomialZ(1)
        elif g.deg < result.deg or not result:
            product = prime
            half = product >> 1
            temp = g * d
            result.init(x if x <= half else x - product
                    for x in temp.coeff)
        elif g.deg == result.deg:
            temp = (g * d - result) * inv_mod_number(product, prime)
            product2 = product * prime
            half = product2 >> 1
            ok = True
            for deg in range(result.deg + 1):
                old = result[deg]
                result[deg] += product * temp[deg]
                if result[deg] > half:
                    result[deg] -= product2
                ok = result[deg] == old
            if ok:
                temp = result.primitive()
                if not m % temp and not n % temp:
                    return temp
            product = product2


@gcd_template.register(PolynomialQ)
def _gcd_polynomial_q(m, n):
    g = lcm_list(x.denominator for x in m.coeff)
    p = PolynomialZ(int(x * g) for x in m.coeff).to_primitive()
    g = lcm_list(x.denominator for x in m.coeff)
    q = PolynomialZ(int(x * g) for x in n.coeff).to_primitive()
    return PolynomialQ(gcd(p, q)).to_monic()


def gcd(m, n):
    return gcd_template(m, n)


def exgcd(m, n):
    klass = type(m)
    a1 = klass()
    a2 = klass(1)
    b1 = klass(1)
    b2 = klass()
    while m:
        (q, m), n = divmod(n, m), m
        a1, a2 = a2, a1 - q * a2
        b1, b2 = b2, b1 - q * b2
    if n:
        leading_coeff = n.leading_coeff
        n /= leading_coeff
        a1 /= leading_coeff
        b1 /= leading_coeff
    return n, a1, b1


def inv_mod(m, n):
    return exgcd(m, n)[1]
